"""
Python语法树解析器
使用Tree-sitter解析Python代码并提取语法信息
"""

import os
from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple, Any
from collections import defaultdict

from tree_sitter_config import (
    TreeSitterConfig, Symbol, FileAnalysis, Position, Range,
    PYTHON_QUERIES, get_node_text, node_to_position, node_to_range
)

try:
    from tree_sitter import Node, Query
except ImportError:
    print("请安装tree-sitter: pip install tree-sitter tree-sitter-python")
    import sys
    sys.exit(1)


class PythonParser:
    """Python代码解析器"""
    
    def __init__(self):
        self.config = TreeSitterConfig()
        if not self.config.is_ready():
            raise RuntimeError("Tree-sitter配置失败")
        
        # 编译查询
        self.queries = {}
        for name, query_str in PYTHON_QUERIES.items():
            try:
                self.queries[name] = Query(self.config.python_language, query_str)
            except Exception as e:
                print(f"警告: 查询 {name} 编译失败: {e}")
    
    def parse_file(self, file_path: str) -> FileAnalysis:
        """
        解析单个Python文件
        
        Args:
            file_path: 文件路径
            
        Returns:
            FileAnalysis: 文件分析结果
        """
        analysis = FileAnalysis(file_path=file_path)
        
        try:
            # 读取文件内容
            with open(file_path, 'rb') as f:
                source_code = f.read()
            
            # 解析生成AST
            tree = self.config.parser.parse(source_code)
            analysis.ast = tree.root_node
            
            # 检查解析错误
            if tree.root_node.has_error:
                analysis.errors.append("语法解析错误")
            
            # 提取符号信息
            analysis.symbols = self._extract_symbols(tree.root_node, source_code, file_path)
            analysis.imports = self._extract_imports(tree.root_node, source_code, file_path)
            
        except Exception as e:
            analysis.errors.append(f"文件解析失败: {str(e)}")
        
        return analysis
    
    def _extract_symbols(self, root_node: Node, source_code: bytes, file_path: str) -> List[Symbol]:
        """提取符号信息"""
        symbols = []
        
        # 提取函数定义
        symbols.extend(self._extract_functions(root_node, source_code, file_path))
        
        # 提取类定义
        symbols.extend(self._extract_classes(root_node, source_code, file_path))
        
        # 提取变量定义
        symbols.extend(self._extract_variables(root_node, source_code, file_path))
        
        return symbols
    
    def _extract_functions(self, root_node: Node, source_code: bytes, file_path: str) -> List[Symbol]:
        """提取函数定义"""
        functions = []
        
        if "function_definitions" not in self.queries:
            return functions
        
        query = self.queries["function_definitions"]
        captures = query.captures(root_node)
        
        current_function = None
        for node, capture_name in captures:
            if capture_name == "function.name":
                function_name = get_node_text(node, source_code)
                # 获取整个函数定义节点
                func_def_node = node.parent
                while func_def_node and func_def_node.type != "function_definition":
                    func_def_node = func_def_node.parent
                
                if func_def_node:
                    current_function = Symbol(
                        name=function_name,
                        type="function",
                        file_path=file_path,
                        range=node_to_range(func_def_node),
                        definition_range=node_to_range(node),
                        scope=self._get_scope(func_def_node)
                    )
                    functions.append(current_function)
        
        return functions
    
    def _extract_classes(self, root_node: Node, source_code: bytes, file_path: str) -> List[Symbol]:
        """提取类定义"""
        classes = []
        
        if "class_definitions" not in self.queries:
            return classes
        
        query = self.queries["class_definitions"]
        captures = query.captures(root_node)
        
        for node, capture_name in captures:
            if capture_name == "class.name":
                class_name = get_node_text(node, source_code)
                # 获取整个类定义节点
                class_def_node = node.parent
                while class_def_node and class_def_node.type != "class_definition":
                    class_def_node = class_def_node.parent
                
                if class_def_node:
                    class_symbol = Symbol(
                        name=class_name,
                        type="class",
                        file_path=file_path,
                        range=node_to_range(class_def_node),
                        definition_range=node_to_range(node),
                        scope=self._get_scope(class_def_node)
                    )
                    
                    # 提取类方法
                    class_symbol.children = self._extract_class_methods(
                        class_def_node, source_code, file_path, class_name
                    )
                    
                    classes.append(class_symbol)
        
        return classes
    
    def _extract_class_methods(self, class_node: Node, source_code: bytes, 
                              file_path: str, class_name: str) -> List[Symbol]:
        """提取类方法"""
        methods = []
        
        def traverse_for_functions(node):
            if node.type == "function_definition":
                name_node = None
                for child in node.children:
                    if child.type == "identifier":
                        name_node = child
                        break
                
                if name_node:
                    method_name = get_node_text(name_node, source_code)
                    method_symbol = Symbol(
                        name=method_name,
                        type="method",
                        file_path=file_path,
                        range=node_to_range(node),
                        definition_range=node_to_range(name_node),
                        scope=f"{class_name}.{method_name}"
                    )
                    methods.append(method_symbol)
            
            for child in node.children:
                traverse_for_functions(child)
        
        traverse_for_functions(class_node)
        return methods
    
    def _extract_variables(self, root_node: Node, source_code: bytes, file_path: str) -> List[Symbol]:
        """提取变量定义"""
        variables = []
        
        if "variable_assignments" not in self.queries:
            return variables
        
        query = self.queries["variable_assignments"]
        captures = query.captures(root_node)
        
        for node, capture_name in captures:
            if capture_name == "variable.name":
                var_name = get_node_text(node, source_code)
                # 获取整个赋值语句节点
                assignment_node = node.parent
                while assignment_node and assignment_node.type != "assignment":
                    assignment_node = assignment_node.parent
                
                if assignment_node:
                    variable = Symbol(
                        name=var_name,
                        type="variable",
                        file_path=file_path,
                        range=node_to_range(assignment_node),
                        definition_range=node_to_range(node),
                        scope=self._get_scope(assignment_node)
                    )
                    variables.append(variable)
        
        return variables
    
    def _extract_imports(self, root_node: Node, source_code: bytes, file_path: str) -> List[Symbol]:
        """提取导入语句"""
        imports = []
        
        if "import_statements" not in self.queries:
            return imports
        
        query = self.queries["import_statements"]
        captures = query.captures(root_node)
        
        for node, capture_name in captures:
            if capture_name == "import.module":
                module_name = get_node_text(node, source_code)
                import_stmt_node = node.parent
                while import_stmt_node and import_stmt_node.type not in ["import_statement", "import_from_statement"]:
                    import_stmt_node = import_stmt_node.parent
                
                if import_stmt_node:
                    import_symbol = Symbol(
                        name=module_name,
                        type="import",
                        file_path=file_path,
                        range=node_to_range(import_stmt_node),
                        definition_range=node_to_range(node),
                        scope="global"
                    )
                    imports.append(import_symbol)
        
        return imports
    
    def _get_scope(self, node: Node) -> str:
        """获取节点的作用域"""
        scope_parts = []
        current = node.parent
        
        while current:
            if current.type == "function_definition":
                # 查找函数名
                for child in current.children:
                    if child.type == "identifier":
                        scope_parts.append(child.text.decode('utf-8'))
                        break
            elif current.type == "class_definition":
                # 查找类名
                for child in current.children:
                    if child.type == "identifier":
                        scope_parts.append(child.text.decode('utf-8'))
                        break
            
            current = current.parent
        
        if not scope_parts:
            return "global"
        
        scope_parts.reverse()
        return ".".join(scope_parts)
    
    def extract_references(self, root_node: Node, source_code: bytes, 
                          symbol_name: str) -> List[Range]:
        """提取符号引用"""
        references = []
        
        if "variable_references" not in self.queries:
            return references
        
        query = self.queries["variable_references"]
        captures = query.captures(root_node)
        
        for node, capture_name in captures:
            if capture_name == "variable.reference":
                ref_name = get_node_text(node, source_code)
                if ref_name == symbol_name:
                    references.append(node_to_range(node))
        
        return references
    
    def extract_function_calls(self, root_node: Node, source_code: bytes) -> List[Tuple[str, Range]]:
        """提取函数调用"""
        calls = []
        
        if "function_calls" not in self.queries:
            return calls
        
        query = self.queries["function_calls"]
        captures = query.captures(root_node)
        
        for node, capture_name in captures:
            if capture_name == "call.function":
                func_name = get_node_text(node, source_code)
                call_range = node_to_range(node)
                calls.append((func_name, call_range))
        
        return calls
    
    def get_ast_structure(self, node: Node, source_code: bytes, depth: int = 0) -> Dict[str, Any]:
        """获取AST结构信息"""
        node_info = {
            "type": node.type,
            "range": f"{node_to_range(node)}",
            "text": get_node_text(node, source_code)[:50] + "..." if len(get_node_text(node, source_code)) > 50 else get_node_text(node, source_code),
            "children": []
        }
        
        if depth < 3:  # 限制深度避免输出过多
            for child in node.children:
                node_info["children"].append(
                    self.get_ast_structure(child, source_code, depth + 1)
                )
        
        return node_info